import numpy as np
import matplotlib.pyplot as plt


class KMeansCustom:
    def __init__(self, init_dataset, n_clusters=2, max_iter=100):
        """
        初始化 KMeansCustom 对象

        :param init_dataset: 初始数据集, 格式为 (N, D), N: 样本数量, D: 特征数量
        :param n_clusters: 聚类数量, default = 2
        :param max_iter: 最大迭代次数, default = 100
        """
        self.init_dataset = init_dataset
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.centroids = None

    def fit(self):
        """
        拟合 KMeans 模型 (聚类)

        :return: KMeansCustom 对象本身
        """
        # 初始化聚类中心为随机选择的样本
        dataset = self.init_dataset
        self.centroids = dataset[np.random.choice(dataset.shape[0], self.n_clusters, replace=False)]

        # 迭代更新聚类中心
        for _ in range(self.max_iter):
            # 计算每个样本与聚类中心的距离
            distances = np.linalg.norm(dataset[:, np.newaxis, :] - self.centroids, axis=2)
            # 根据距离将样本分配到最近的聚类中心
            labels = np.argmin(distances, axis=1)
            # 更新聚类中心为每个簇的平均值
            new_centroids = np.array([dataset[labels == k].mean(axis=0) for k in range(self.n_clusters)])
            if np.all(self.centroids == new_centroids):
                break
            self.centroids = new_centroids

        return self

    def predict(self, new_dataset):
        """
        预测新样本所属的簇

        :param new_dataset: 新的样本数据, 格式为 (N, D), N: 样本数量, D: 特征数量
        :return: 新样本的聚类标签
        """
        distances = np.linalg.norm(new_dataset[:, np.newaxis, :] - self.centroids, axis=2)
        labels = np.argmin(distances, axis=1)
        return labels

    def plot_cluster_result(self, title):
        """
        绘制聚类结果的散点图

        :param title: 图片标题
        """
        # 获取聚类结果
        dataset = self.init_dataset
        cluster_labels = self.predict(dataset)
        # 绘制分类结果
        plt.figure(figsize=(8, 6))
        plt.scatter(self.centroids[:, 0], self.centroids[:, 1], marker='s', c='red', linewidths=5, label='Centroids')
        plt.scatter(data[:, 0], data[:, 1], c=cluster_labels, cmap='viridis', edgecolors='k', alpha=0.7)
        plt.xlabel('Feature 1')
        plt.ylabel('Feature 2')
        plt.title(title)
        plt.legend()
        plt.grid(True)
        plt.show()


# 创建数据集
data = np.array([[170, 70], [178, 75], [100, 100], [120, 40], [10, 0.1]])

if __name__ == '__main__':
    for n in range(2, 4):
        kmeans = KMeansCustom(init_dataset=data, n_clusters=n)
        kmeans.fit()
        kmeans.plot_cluster_result(f'KMeans Clustering (k={n})')
